# -*- coding: utf-8 -*-
"""
Created on Thu Aug 29 11:26:54 2013

@author: jkwong
"""

import base64
import getpass
import os
import socket
import sys
import traceback
import paramiko
import serial
import numpy as np

def ReadZspec(fullfilenameList):
    """ Read Zspec data """
    
    # Check if the input is a string (single file)
    if isinstance(fullfilenameList, str):
        temp = []
        temp.append(fullfilenameList)
        fullfilenameList = temp
    
    dat = list()
    for ii in range(len(fullfilenameList)):
        temp = np.genfromtxt(fullfilenameList[ii], \
                                 delimiter=',', \
                                 skip_header = 0, \
                                 skip_footer = 0, \
                                 dtype = 'uint32')
        if (temp.shape[0] == 257):  # shave off the table labels if present
            temp = temp[1:,:]
        if (temp.shape[1] == 138):
            temp = temp[:,1:-1]
        if (temp.shape[1] == 137): # something with newline causeing it to think there is one more detetor
            temp = temp[:,1:]
        dat.append(temp)
    return dat


def CheckGain(*arg):
    """Given a filename, acquisition time and threshold, it gives the detector number
        gain shift, and answer as to whether or not the gain shift has exceeded the 
        threshold.
        8/28/2013, John Kwong
        8/29/2013, Added option to print all or 
    """
    #parse the arguments
    fullFilename, acqTime, rateRange, rateLook, binBaseline, badDetectors, threshold = arg[0:7]
    # convert the threshold string to float
    threshold = float(threshold)
    # convert acquisition time
    acqTime = float(acqTime)
    if len(arg) < 4:
        option = 'all'
    else:
        option = arg[3]
    
    # check if the option is valid
    if option not in ['badgain', 'gooddet', 'all', 'badgaingooddet']:
        print("Invalid Option.")
        return

#    # Baseline values
#    # Fit exponential to this region
#    rateRange = np.array([0.5, 2.0])
#
#    # bin value of the fitted exponential at count rate
#    rateLook = 1.0    
#    # should be at this value
#    binBaseline = 100.0
#    # the bin value have to be between these values
    binBounds = [20, 200]
#    
#    # index starts at 1
#    badDetectors = np.array([1,2,3,4,5,6,7,8,20,26,31,33,34,38,39,40,44,53,56,62,68,76,80,125,126,127,128,129,130,131,132,133,134,135,136])
        
    # Read the file
    dat = ReadZspec(fullFilename)
    
    # Shift among PMTs - OLD STUFF 8/28/2013
#    binMeanCC = np.zeros(dat[0].shape[1])
#    gainShift = np.zeros(dat[0].shape[1])
    
#    print dat[0].shape
    binArray = np.arange(dat[0].shape[0])
    
    # fit an exponential
#    pfitAll = np.zeros((dat[0].shape[1], 2))
    binMeanAll = np.zeros(dat[0].shape[1])

    for i in xrange(dat[0].shape[1]):
        cutt = (binArray > binBounds[0]) & (binArray < binBounds[1]) & (dat[0][:,i]/acqTime > rateRange[0]) & (dat[0][:,i]/acqTime < rateRange[1])
        if any(cutt):
            temp = np.polyfit(binArray[cutt], np.log(dat[0][cutt,i]/acqTime),1)
            pfit0 = [temp[0], np.exp(temp[1])]
            binMean = (1.0/pfit0[0]) * np.log(rateLook/pfit0[1])
        else:
            binMean = 0.0
        binMeanAll[i] = binMean
#        print i+1, binMean
    gainShift = binMeanAll / binBaseline

#    # OLD STUFF - 8/28/2013, JK
#    binMat = np.tile(np.arange(dat[0].shape[0]),(dat[0].shape[1],1)).T
#    cutt = (binMat > binBounds[0]) & (binMat < binBounds[1]) & (dat[0]/acqTime > rateRange[0]) & (dat[0]/acqTime < rateRange[1])
#
#    # calculate the bin value for all detectors
#    binMeanCC = (binMat * dat[0] * cutt).sum(axis = 0).astype(float) / (dat[0] * cutt).sum(axis = 0).astype(float)
#
#    gainShift = binMeanCC / np.tile(binBaseline, dat[0].shape[1])    
#    # If gain is okay, then score[i] = true


    score = np.abs(1 - gainShift) <= threshold

    # display results
    print("Target value at rate %3.3f is %3.3f" %(rateLook, binBaseline))
    print("Gain Shift Threshold =  %3.3f" % threshold)
    
    print("%s\t%s\t%s" %("Det#", "Shift", "Good?"))
    
    # detector # = index+1
    for (index, gs) in enumerate(gainShift):
        # Only show for not bad detectors
        if option == 'all':
                print("%d\t%3.3f\t%r" %(index+1, gs, score[index]))
        elif option == 'gooddet':
            if (index+1) not in badDetectors:
                print("%d\t%3.3f\t%r" %(index+1, gs, score[index]))
        elif option == 'badgain':
            if ~score[index]:
                print("%d\t%3.3f\t%r" %(index+1, gs, score[index]))
        elif option == 'badgaingooddet':
            if ~score[index] and (index+1) not in badDetectors:
                print("%d\t%3.3f\t%r" %(index+1, gs, score[index]))
    return()
    
def GetZspecFile(sshObj, pathHost = r'/home/dastest/vxWorks/', pathLocal = r'C:\Users\jkwong\Documents\Work\PBAR\data'):
    """Gets latest or specific zspec file from host"""

    # include the final slashes
    pathMar =  pathHost + 'data/'
    pathAnalysis =  pathHost + r'analysis/'
    
#    sshObj = PBAR_HighVoltageCalibrator.SecureConnect('pbarsite.no-ip.org', 'dastest', 'School bundle plate chimney 10', 22)
    sshObj.ConnectSSHClient()
    sshObj.ConnectSFTP()
    
    # Get a list of the 20 most recent mar files
    stdin, stdout, stderr = sshObj.client.exec_command('ls -lt %s*.mar | head -n20' % pathMar)
    lines = []
    for s in stdout:
        lines.append(s[:-1])
        
    # Get the latest mar file
    # Start from most recent file. If not mar file, then skip to next file. if mar file then accept as the latest file
    # Also parse file prefix and file number
    filename, fileNumber, filePrefix = None, None, None
    for (index, line) in enumerate(lines):
        # find start index of file name
        if len(line) > 3:
            if line[-3:] == 'mar':
                startIndex = line.find('data/')+5
                filename = line[startIndex:]
                fileNumber = int(filename[2:4])
                filePrefix = filename[0:2]
                break
    
    # Create associated csv if file has been found
    if filename != None:
        # string looks like this: python /home/dastest/vxWorks/analysis/PBAR_Zspec_MarToCSV.py ea 28 28
        command = 'python %sPBAR_Zspec_MarToCSV.py %s %d %d' %(pathAnalysis, filePrefix, fileNumber, fileNumber)
        stdin, stdout, stderr = sshObj.client.exec_command(command)
    else:
        print('No mar file found in %s. Quitting.' %pathMar)
        return
    
    # Copy the csv file to local directory
    filenameCSV = filename.replace('mar', 'csv')
    fullFilenameServer = '%s%s' %(pathAnalysis, filenameCSV)
    fullFilenameLocal = os.path.join(pathLocal, filenameCSV)
    sshObj.sftp.get(fullFilenameServer, fullFilenameLocal)
    
#    # Disconnect clients
#    sshObj.DisconnectSFTP()
#    sshObj.DisconnectSSHClient()
    
    
class HighVoltageSystem:
    def __init__(self):
        # list of high voltage devices
        self.hvDevices = {}
        self.map
    def ReadVoltages(self, detectorList):
        1+1
    def SetVoltages(self):
        1+1

class HighVoltageDevice:
    def __init__(self, name, COM):
        self.name = name
        self.COM = COM  # COM port number

        # things to be set    
        self.serial = None
    
    def Connnect(self):
        1+1
    def Disconnect(self):
        self.serial.close()
    def ReadVoltages(self, channelList):
        1+1
    def SetVoltages(self, channelList):
        1+1
    
class SecureConnect:
    def __init__(self, hostname, username, password, port):
        self.hostname = hostname
        self.username = username
        self.password = password # convert to keyless later
        self.port = port
        
        # keys
        self.hostkeytype = None
        self.hostkey = None
        self.host_keys = None
        
        # containers for stfp and ssh connections
        self.sftp = None
        self.client = None
        self.transport = None
        # setup logging
        paramiko.util.log_to_file('demo_sftp.log')

    def GetHostKeys(self):    
        # get host key, if we know one

        try:
            self.host_keys = paramiko.util.load_host_keys(os.path.expanduser('~/.ssh/known_hosts'))
        except IOError:
            try:
                # try ~/ssh/ too, because windows can't have a folder named ~/.ssh/
                self.host_keys = paramiko.util.load_host_keys(os.path.expanduser('~/ssh/known_hosts'))
            except IOError:
                print '*** Unable to open host keys file'
                self.host_keys = {}
        
        if self.host_keys.has_key(self.hostname):
            self.hostkeytype = self.host_keys[self.hostname].keys()[0]
            self.hostkey = self.host_keys[self.hostname][self.hostkeytype]
            print 'Using host key of type %s' % self.hostkeytype
    def ConnectSFTP(self):

        # make sftp object
        try:
            self.transport = paramiko.Transport((self.hostname, self.port))
            self.transport.connect(username=self.username, password=self.password, hostkey=self.hostkey)
            self.sftp = paramiko.SFTPClient.from_transport(self.transport)
        except:
            print("Could not connect to server.")
    def DisconnectSFTP(self):
        self.transport.close()
        self.sftp = None
    
    def ConnectSSHClient(self):
        self.client = paramiko.SSHClient()
        self.client.load_system_host_keys()
        self.client.connect(hostname = self.hostname, username = self.username, password = self.password)

    def DisconnectSSHClient(self):
        self.client.close()
        
    def SSHClietListFiles(self, path):
        """List files in time order starting with most recent"""
        (stdin, stdout, stderr) = self.client.exec_command('ls -ltr ' + path)
        return(stdout.readlines())


#class HighVoltageController:
#    def __init__(self):
#        highVoltageInterfaces = []
#
#class HighVoltageInterface:

#class PBAR_VoltageCalculator:

    
